from langdev_helper.llm.lcex import llm_lcex as model


from typing import List, Tuple
from typing_extensions import Annotated
from langchain_core.documents import Document
from langchain_core.messages import ToolMessage
from langchain_core.tools import tool
from langgraph.prebuilt import InjectedState
from pydantic import BaseModel

@tool(parse_docstring=True, response_format="content_and_artifact")
def get_context(question: List[str]) -> Tuple[str, List[Document]]:
    """获取关于问题的上下文。

    Args:
        question: 用户问题
    """
    print('\n--tool--\n', 'get_context called')
    # 返回常量虚拟输出
    docs = [
        Document(
            "FooBar 公司刚刚筹集了10亿美元！",
            metadata={"source": "twitter"},
        ),
        Document(
            "FooBar 公司现在只雇用 AI", metadata={"source": "twitter"}
        ),
        Document(
            "FooBar 公司成立于2019年", metadata={"source": "wikipedia"}
        ),
        Document(
            "FooBar 公司制造友好的机器人", metadata={"source": "wikipedia"}
        ),
    ]
    return "\n\n".join(doc.page_content for doc in docs), docs

@tool(parse_docstring=True, response_format="content_and_artifact")
def cite_context_sources(claim: str, state: Annotated[dict, InjectedState]) -> Tuple[str, List[Document]]:
    """引用支持某一声明的来源。

    Args:
        claim: 声明内容
    """
    print('\n--tool--\n', 'cite_context_sources called')
    docs = []
    # 从过去的 ToolMessages 中获取可能引用的文档
    for msg in state["messages"]:
        if isinstance(msg, ToolMessage) and msg.name == "get_context":
            docs.extend(msg.artifact)

    class Cite(BaseModel):
        """返回支持声明的文档索引"""
        indexes: List[int]

    structured_model = model.with_structured_output(Cite)
    system = f"以下哪些文档最能证明该声明:\n\n{claim}"
    context = "\n\n".join(
        f"文档 {i}:\n" + doc.page_content for i, doc in enumerate(docs)
    )
    citation = structured_model.invoke([("system", system), ("human", context)])
    cited_docs = [docs[i] for i in citation.indexes]
    sources = ", ".join(doc.metadata["source"] for doc in cited_docs)
    return sources, cited_docs